import json
from typing import Dict, Tuple
from volcengine.visual.VisualService import VisualService
from exception import AppException
from common.enums.apply import DrawsEnum
from common.utils.tools import ToolsUtil
from common.utils.config import ConfigUtil
from common.cache.pool_cache import KeyPoolCache


class DrawsMjServer:
    """ MJ绘画 """
    def __init__(self, params: Dict[str, any]) -> None:
        # 绘画渠道
        self.channel = "zsy"

        # 基础参数
        self.prompt: str = params.get("prompt", "")               # 提示词
        self.action: str = str(params.get("action", ""))          # 操作类型: [generate, upsample, variation]
        self.model: str = str(params.get("model", ""))            # 生成模式: [mj, niji]
        self.scale: str = str(params.get("scale", ""))            # 图片比例: [1:1, 1:2, ...]
        self.version: str = str(params.get("version", ""))        # 版本参数: [6.0, 5.2, ...]
        self.image_id: str = str(params.get("image_id", ""))      # 图片的ID
        self.image_ref: str = str(params.get("image_ref", ""))    # 参考图片
        self.callback: str = params.get("callback_url", "")       # 回调地址

        # 其它参数
        others = params.get("other_params") or {}
        self.mode: str = str(others.get("mode", ""))              # 生成模式: [fast, relax, turbo]
        self.style: str = str(others.get("style", ""))            # 风格样式: [default=动漫, cute=可爱, ...]
        self.no_content: str = str(others.get("no_content", ""))  # 排除内容: [去除图片的某个东西]
        self.other_params = {
            "mode": self.mode,
            "style": self.style,
            "no_content": self.no_content
        }

        # 高级参数
        self.complex_params = {
            # 种子参数
            "seed": params.get("complex_params", {}).get("seed", ""),
            # 图片权重: 默认1,[0.5-2],数值越高越接近参考图,反之偏向描述词
            "iw": float(params.get("complex_params", {}).get("iw", "1")),
            # 图片质量: 默认1,[0.25,1],图片细节刻画,越小出图越快
            "q": float(params.get("complex_params", {}).get("q", "1")),
            # 图片模糊: 默认100,[0-100],参数越高,风格化细节越多,生成时间越久
            "s": int(params.get("complex_params", {}).get("s", "100")),
            # 图片随机: 默认0,[0-100],数值越大随机性越高,较低则还是按照描述词走
            "c": int(params.get("complex_params", {}).get("c", "0"))
        }

        # 验证参数
        self._check_params()

        # 构建指令
        self._make_instruct()

    async def imagine(self):
        """ 生成图片 """
        config = await ConfigUtil.get("draws", "mj") or {}
        self.channel = config.get("channel", "")

        if self.channel == "zsy":
            return await self.post_zsy()
        elif self.channel == "acedata":
            return await self.post_aceData()
        else:
            raise AppException("无可用绘画通道,请管理员检查配置信息。")

    async def post_zsy(self):
        """
        知识云MJ

        Document:
            https://data.zhishuyun.com/documents

        Author:
            zero
        """
        url: str = "https://api.zhishuyun.com/midjourney/imagine"
        if self.mode == "relax":
            url += "/relax"
        elif self.mode == "turbo":
            url += "/turbo"

        token: str = await self.get_apikey()
        response = await ToolsUtil.post_curl(
            url=f"{url}?token=" + token,
            headers={
                "Accept": "application/json",
                "Content-Type": "application/json"
            },
            data={
                "action": self.action,
                "prompt": self.prompt,
                "image_id": self.image_id,
                "callback_url": self.callback,
                "translation": True,
                "timeout": 100
            })
        return self._handle_response(response)

    async def post_aceData(self):
        """
        AceData

        Document:

        """
        token: str = await self.get_apikey()
        response = await ToolsUtil.post_curl(
            url="https://api.acedata.cloud/midjourney/imagine",
            headers={
                "Accept": "application/json",
                "Content-Type": "application/json",
                "Authorization": "Bearer " + token
            },
            data={
                "mode": self.mode,
                "action": self.action,
                "prompt": self.prompt,
                "image_id": self.image_id,
                "mask": self.image_ref,
                "callback_url": self.callback,
                "translation": True,
                "timeout": 100
            })
        return self._handle_response(response)

    async def get_apikey(self) -> str:
        """ 获取密钥 """
        if self.channel == "acedata":
            marks = f"mj-{self.channel}"
        else:
            marks = f"mj-{self.channel}-{self.mode}"
        apikey = await KeyPoolCache(scene="draws", channel=marks).get_key() or {}
        if not apikey or not apikey.get("api_key"):
            raise AppException("尚未配置密钥,请与管理员联系")
        return apikey["api_key"]

    def _make_instruct(self):
        """ 构建指令 """
        # 参考图片
        if self.image_ref:
            self.prompt = self.image_ref + " " + self.prompt
        # 图片比例
        if self.scale:
            self.prompt += f" --ar {self.scale}"
        # 排除内容
        if self.no_content:
            self.prompt += f" --no {self.no_content}"
        # 版本参数
        if self.version:
            if self.version.startswith("--v") or self.version.startswith("--niji"):
                self.prompt += f" {self.version}"
                if self.style and self.style != "default":
                    self.prompt += f" --style {self.style}"
        # 高级参数
        if self.complex_params:
            complex_params = ""
            for k, v in self.complex_params.items():
                if not v:
                    continue
                complex_params += f" --{k} {v}"
            if complex_params:
                self.prompt += complex_params

    def _check_params(self):
        """ 验证参数 """

        if self.action not in ["generate", "upscale1", "upscale2", "upscale3", "upscale4",
                               "rerzoll", "variation1", "variation2", "variation3", "variation4"]:
            raise AppException("操作类型错误了")

        if self.action != "generate" and not self.image_id:
            raise AppException("请指定操作图片")

        if self.mode not in ["fast", "relax", "turbo"]:
            raise AppException("生成模式选择异常")

        if self.model not in ["mj", "niji"]:
            raise AppException("生成模型选择异常")

        if self.scale not in ["1:1", "1:2", "4:3", "9:16", "3:4", "16:9", "3:2", "2:3"]:
            raise AppException("不支持的图片比例")

        mj_version = (DrawsEnum.mj_config("mj_version") or {}).get("mj")
        if not mj_version and self.model == "mj" and self.version not in mj_version:
            raise AppException("模型版本选择异常")

        if self.model == "niji":
            mj_style = DrawsEnum.mj_config("mj_style") or {}
            niji_version = (DrawsEnum.mj_config("mj_version") or {}).get("niji")
            if self.version not in niji_version:
                raise AppException("模型版本选择异常")
            if self.style not in mj_style.keys():
                raise AppException("生成风格选择异常")

    @classmethod
    def _handle_response(cls, response):
        if response.get("code") == "invalid_token":
            raise Exception("令牌无效, 请确保您的密钥正确。")
        elif response.get("code") == "token_mismatched":
            raise Exception("令牌不匹配, 请确保您的密钥正确。")
        elif response.get("code") == "used_up":
            raise Exception("API套餐已经用完，请管理员到平台充值。")
        elif response.get("code") == "unknown":
            raise Exception("未知错误，需要联系API平台系统管理员。")
        elif response.get("code") == "unverified":
            raise Exception("未完成实名认证, 需管理员到API平台进行实名认证。")
        elif response.get("code") == "forbidden":
            raise Exception(response.get("detail", "forbidden"))
        return response


class DrawsDoubaoServer:
    """ 豆包绘画 """
    def __init__(self, params: Dict[str, any] = None) -> None:
        self.visual_service = None
        self.params = params
        self.complex_params = {}

    async def serv(self):
        apikey = await self.get_apikey()
        self.visual_service = VisualService()
        self.visual_service.set_ak(apikey.get("api_key"))
        self.visual_service.set_sk(apikey.get("api_secret"))
        return self.visual_service

    @classmethod
    async def get_apikey(cls) -> dict:
        """ 获取密钥 """
        apikey = await KeyPoolCache(scene="draws", channel="doubao").get_key() or {}
        if not apikey or not apikey.get("api_key"):
            raise AppException("尚未配置密钥,请与管理员联系")
        return apikey

    async def imagine(self):
        if self.params.get("model") == "high_aes_general_v21_L":
            return await self.text_to_image()
        elif self.params.get("model") == "high_aes":
            return await self.anime_to_image()
        elif self.params.get("model") == "i2i_xl_sft":
            return await self.image_to_image()
        else:
            raise AppException(f"绘画模型异常: {self.params.get('model')}")

    async def text_to_image(self):
        """
        文生图 (通用2.1-文生图)

        Document:
            https://www.volcengine.com/docs/6791/1366783

        Author:
            zero
        """
        await self.serv()

        try:
            width, height = self._scale_to_wh(self.params.get("scale", "512x512"))
            complex_params = self.params.get("complex_params", {})
            self.complex_params = {
                "seed": int(complex_params.get("seed") or -1),
                "scale": float(complex_params.get("scale") or 3.5),
                "ddim_steps": int(complex_params.get("ddim_steps") or 25)
            }

            resp = self.visual_service.cv_sync2async_submit_task({
                # 算法名称
                "req_key": "high_aes_general_v21_L",
                # 模型版本
                "model_version": "general_v2.1_L",
                # 美感配置
                "req_schedule_conf": "general_v20_9B_pe",
                # 提示词值
                "prompt": str(self.params.get("prompt")),
                # 随机种子: 默认值: [-1]
                "seed": self.complex_params["seed"],
                # 影响文本: 默认3.5, [1~10]
                "scale": self.complex_params["scale"],
                # 生成步数: 默认25, [1~200]
                "ddim_steps": self.complex_params["ddim_steps"],
                # 图像宽度
                "width": width,
                # 图像高度
                "height": height,
                # 文本扩写: [针对prompt优化]
                "use_pre_llm": True,
                # 超分功能: [True=文生图+AIGC超分, False=文生图]
                "use_sr": True
            })
            return self._handle_response(resp)
        except Exception as e:
            self._handle_exception(e)

    async def image_to_image(self):
        """
        图生图 (通用XL pro-图生图)

        Document:
            https://www.volcengine.com/docs/6791/1330203

        Author:
            zero
        """
        await self.serv()

        try:
            complex_params = self.params.get("complex_params", {})
            self.complex_params = {
                "seed": int(complex_params.get("seed") or -1),
                "scale": float(complex_params.get("scale") or 7.0),
                "ddim_steps": int(complex_params.get("ddim_steps") or 20),
            }

            resp = self.visual_service.cv_sync2async_submit_task({
                # 算法名称
                "req_key": "i2i_xl_sft",
                # 提示词值
                "prompt": self.params.get("prompt"),
                # 随机种子: 默认-1
                "seed": self.complex_params["seed"],
                # 影响文本: 默认7.0, 范围[1~30]
                "scale": self.complex_params["scale"],
                # 生成步数: 默认20, 范围[1~50]
                "ddim_steps": self.complex_params["ddim_steps"],
                # 参考图片
                "image_urls": [self.params.get("image_ref")] if self.params.get("image_ref", None) else []
            })
            return self._handle_response(resp)
        except Exception as e:
            self._handle_exception(e)

    async def anime_to_image(self):
        """
        动漫图 (动漫1.3.X-文生图/图生图)

        Document:
            https://www.volcengine.com/docs/6791/1213131

        Author:
            zero
        """
        await self.serv()

        try:
            width, height = self._scale_to_wh(self.params.get("scale", "512x512"))
            complex_params = self.params.get("complex_params", {})
            self.complex_params = {
                "seed": int(complex_params.get("seed") or -1),
                "scale": float(complex_params.get("scale", 7.0)),
                "strength": float(complex_params.get("strength", 0.7)),
                "ddim_steps": int(complex_params.get("ddim_steps", 25))
            }

            resp = self.visual_service.cv_sync2async_submit_task({
                # 算法名称
                "req_key": "high_aes",
                # 模型版本
                "model_version": "anime_v1.3.1",
                # 提示词值
                "prompt": self.params.get("prompt"),
                # 随机种子: 默认-1
                "seed": self.complex_params["seed"],
                # 影响文本: 默认7, [1~30]
                "scale": self.complex_params["scale"],
                # 接近原图, 默认0.7, [0~1.0]
                "strength": self.complex_params["strength"],
                # 生成步数: 默认25, [1~200]
                "ddim_steps": self.complex_params["ddim_steps"],
                # 参考图片
                "image_urls": [self.params.get("image_ref")] if self.params.get("image_ref", None) else [],
                # 图像宽度
                "width": width,
                # 图像高度
                "height": height
            })
            return self._handle_response(resp)
        except Exception as e:
            self._handle_exception(e)

    async def fetch(self, req_key: str, task_id: str):
        """
        获取生成结果

        Args:
            req_key (str): key
            task_id (str): 任务ID

        Returns:
            dict

        Author:
            zero
        """
        try:
            await self.serv()
            body = self.visual_service.cv_sync2async_get_result({
                "req_key": req_key,
                "task_id": task_id,
                "req_json": json.dumps({
                    "return_url": True
                })
            })
            return self._handle_response(body)
        except Exception as e:
            raise AppException(str(e))

    @classmethod
    def _scale_to_wh(cls, scale: str) -> Tuple[int, int]:
        """
        比例转宽高

        Args:
            scale (str): 512x512

        Returns:
            tuple: (width, height)

        Author:
            zero
        """
        if not scale:
            return 512, 512
        scales = scale.split("x")
        if len(scales) >= 2:
            return int(scales[0]), int(scales[1])
        return 512, 512

    @classmethod
    def _handle_response(cls, response):
        """
        处理响应结果

        Args:
            response (dict): {"code": 10000, "data": {"task_id": "7468667968", "message": "Success"}}

        Returns:
            dict:
                {"task_id": '7468667968814088203'}
                {"image_urls": ["https://xxx.png"]}
                {'binary_data_base64': [], 'image_urls': None, 'resp_data': '', 'status': 'in_queue'}

        Author:
            zero
        """
        if not response or response.get("status") != 10000 or not response.get("data"):
            code: str = str(response.get("code"))
            message: str = str(response.get("message"))
            raise AppException(f"{code}:{message}")

        return response.get("data")

    @classmethod
    def _handle_exception(cls, e):
        try:
            except_string = str(e).strip("b'").strip("'")
            decoded_string = except_string.encode("utf-8")
            decoded_string = decoded_string.decode("utf-8")
            error = json.loads(decoded_string)
            wrong = error["ResponseMetadata"]["Error"]
            code = wrong.get("Code")
            message = wrong.get("Message")
            if code == "InvalidAccessKey":
                raise Exception("令牌无效, 请确保您的密钥正确。")
            elif code == "FlowLimitExceeded":
                raise Exception("请求过于频繁, 超出了限速(QPS), 请稍后重试。")
        except Exception as e:
            raise Exception(f"1: {str(e)}")
        raise Exception(f"{code}: {message}")


class DrawsDalleServer:
    """ Dalle绘画 """

    def __init__(self, params: Dict[str, any] = None):
        self.params = params
        self.baseUrl = "https://api.openai.com/v1"
        other_params = self.params.get("other_params", {})
        self.other_params = {
            "quality": other_params.get("quality") or "standard",
            "style": other_params.get("style") or "vivid"
        }

    async def imagine(self):
        # 获取配置
        config = await ConfigUtil.get("draws", "dalle") or {}
        if config.get("agency_api"):
            self.baseUrl = config["agency_api"].strip("/").strip()

        # 获取密钥
        token: str = await self.get_apikey()

        # 发起请求
        try:
            return await ToolsUtil.post_curl(
                url=f"{self.baseUrl}/images/generations",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": "Bearer " + token
                },
                timeout=300,
                data={
                    "n": 1,
                    "model": self.params.get("model") or "dall-e-3",
                    "prompt": self.params.get("prompt"),
                    "size": self.params.get("scale") or "1:1",
                    "quality": self.other_params["quality"],
                    "style": self.other_params["style"],
                    "response_format": "b64_json"
                })
        except Exception as e:
            raise AppException(str(e))

    @classmethod
    async def get_apikey(cls) -> str:
        """ 获取密钥 """
        apikey = await KeyPoolCache(scene="draws", channel="dalle").get_key() or {}
        if not apikey or not apikey.get("api_key"):
            raise AppException("尚未配置密钥,请与管理员联系")
        return apikey["api_key"]


class DrawsSdServer:
    """ SD绘画 """

    def __init__(self):
        # 接收参数
        # 获取密钥
        pass

    async def imagine(self):
        data = {
            "prompt": "",
            "negative_prompt": "",
            "seed": "",
            "sampler_name": "",
            "steps": "",
            "cfg_scale": "",
            "batch_size": "",
            "width": "",
            "height": ""
        }

        # 文生图
        url = "/sdapi/v1/txt2img"

        # 图生图
        url = "sdapi/v1/img2img"

        # 发起请求
        await ToolsUtil.post_curl(
            url=url,
            headers={
                "Accept": "application/json",
                "Content-Type": "application/json"
            },
            data=data
        )
